import os
import json
from typing import Optional, Callable
from src.embeddings.transformer_embedding import TransformerEmbedding
from src.llm.base import BaseLLM
from src.server.config import get_settings
from src.vectorstore.faiss_store import FAISSStore
from src.retriever.vector_retriever import VectorRetriever
from src.rag.simple_rag import SimpleRAG
from src.loaders.document_loader import DocumentLoader

PROCESSED_FILES_RECORD = "data/vector_store/processed_files.json"

async def initialize_rag_system(
    progress_callback: Optional[Callable] = None
) -> SimpleRAG:
    """初始化RAG系统"""
    try:
        print("正在初始化RAG系统...")
        
        # 初始化组件
        embedding_model = TransformerEmbedding()
        vectorstore = FAISSStore(embedding_model)
        retriever = VectorRetriever(vectorstore=vectorstore, k=4)
        
        # 从配置中获取LLM类型
        settings = get_settings()
        llm_model = settings.LLM_MODEL  # 确保配置中使用正确的模型名称
        
        try:
            llm = BaseLLM.get_llm(llm_model)
        except ValueError as e:
            print(f"初始化LLM时出错: {str(e)}")
            raise
        
        # 初始化RAG
        rag_system = SimpleRAG(
            embedding_model=embedding_model,
            vectorstore=vectorstore,
            retriever=retriever,
            llm=llm,
            k=4
        )
        
        # 加载新增的文档
        await load_new_documents(rag_system, progress_callback)
        
        return rag_system
        
    except Exception as e:
        print(f"初始化RAG系统时出错: {str(e)}")
        raise

def get_processed_files() -> set:
    """获取已处理的文件列表"""
    try:
        if os.path.exists(PROCESSED_FILES_RECORD):
            with open(PROCESSED_FILES_RECORD, 'r', encoding='utf-8') as f:
                return set(json.load(f))
    except Exception as e:
        print(f"读取处理记录时出错: {str(e)}")
    return set()

def save_processed_files(processed_files: set) -> None:
    """保存已处理的文件列表"""
    try:
        os.makedirs(os.path.dirname(PROCESSED_FILES_RECORD), exist_ok=True)
        with open(PROCESSED_FILES_RECORD, 'w', encoding='utf-8') as f:
            json.dump(list(processed_files), f, ensure_ascii=False, indent=2)
    except Exception as e:
        print(f"保存处理记录时出错: {str(e)}")

async def load_new_documents(
    rag_system: SimpleRAG,
    progress_callback: Optional[Callable] = None
) -> None:
    """加载新增的文档"""
    try:
        print("\n=== 检查新增文档 ===")
        
        # 获取已处理的文件列表
        processed_files = get_processed_files()
        
        # 设置文档目录
        docs_dir = "data/documents"
        if not os.path.exists(docs_dir):
            print(f"创建文档目录: {docs_dir}")
            os.makedirs(docs_dir)
            return
            
        # 加载新文档
        loader = DocumentLoader()
        new_documents = []
        new_metadata = []
        new_files = set()
        
        # 遍历文档目录
        for filename in os.listdir(docs_dir):
            if not filename.endswith(('.txt', '.pdf', '.docx')):
                continue
                
            file_path = os.path.join(docs_dir, filename)
            file_info = f"{filename}:{os.path.getmtime(file_path)}"  # 文件名和修改时间
            
            # 检查文件是否已处理
            if file_info in processed_files:
                continue
                
            print(f"发现新文件: {filename}")
            doc_content = loader.load_document(file_path)
            if doc_content:
                new_documents.append(doc_content)
                new_metadata.append({
                    'source': filename,
                    'path': file_path,
                    'type': filename.split('.')[-1]
                })
                new_files.add(file_info)
        
        if new_documents:
            print(f"\n=== 添加文档开始 ===")
            print(f"文档数量: {len(new_documents)}")
            
            if progress_callback:
                await progress_callback(f"正在添加 {len(new_documents)} 个文档...")
                
            # 添加到RAG系统
            rag_system.add_documents(new_documents, new_metadata)
            
            # 更新处理记录
            processed_files.update(new_files)
            save_processed_files(processed_files)
            
            print("=== 添加文档完成 ===")
        else:
            print("没有新增文档需要处理")
            
    except Exception as e:
        print(f"加载新增文档时出错: {str(e)}")
        raise